package xyz.ylx.apirotation.service.impl;

import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import xyz.ylx.apirotation.config.GeminiApiKeyConfig;
import xyz.ylx.apirotation.service.AIModelService;
import xyz.ylx.apirotation.exception.custom.ApiKeyException;
import xyz.ylx.apirotation.common.util.ApiKeyUtil;

import java.util.*;

import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@Service
@AllArgsConstructor
public class GeminiProxyServiceImpl implements AIModelService {
    
    private final GeminiApiKeyConfig apiKeyConfig;
    private final WebClient geminiWebClient;
    
    @Override
    public Object handleChatCompletions(Map<String, Object> requestBody, Map<String, String> headers, boolean isStream) {
        if (isStream) {
            return handleOpenAIStreamChatCompletions(requestBody, headers);
        } else {
            return Mono.error(new UnsupportedOperationException("暂不支持Gemini的非流式响应"));
        }
    }
    
    @Override
    public boolean supportsModel(String modelName) {
        if (!StringUtils.hasLength(modelName)) {
            return false;
        }
        
        // 根据前缀判断是否支持模型
        return modelName.startsWith("gemini-") || 
               modelName.startsWith("gemini/") || 
               modelName.startsWith("gemini_");
    }
    
    public Flux<String> handleOpenAIStreamChatCompletions(
            Map<String, Object> requestBody, 
            Map<String, String> headers) {
        
        List<String> keys = apiKeyConfig.getKeys();
        if (CollectionUtils.isEmpty(keys)) {
            return Flux.error(new ApiKeyException("未配置Gemini API keys，请在配置文件中设置gemini.api.keys"));
        }
        
        Optional<String> apiKeyOpt = ApiKeyUtil.getNextApiKey(keys);
        if (!apiKeyOpt.isPresent()) {
            return Flux.error(new ApiKeyException("无法获取Gemini API key"));
        }
        String apiKey = apiKeyOpt.get();
        
        log.info("开始请求Gemini API，参数：{}", requestBody);
        
        return geminiWebClient
            .post()
            .uri("/chat/completions")
            .header(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey)
            .headers(h -> {
                // 复制一些重要的客户端请求头
                if (headers.containsKey(HttpHeaders.USER_AGENT)) {
                    h.set(HttpHeaders.USER_AGENT, headers.get(HttpHeaders.USER_AGENT));
                }
            })
            .bodyValue(requestBody)
            .retrieve()
            .bodyToFlux(String.class)
            .doOnSubscribe(s -> log.debug("开始订阅流"))
            .doOnNext(line -> {
                log.debug("收到响应: {}", line);
            })
            .doOnError(error -> {
                log.error("请求错误: {}", error.getMessage(), error);
            })
            .doOnComplete(() -> {
                log.debug("流请求完成");
            });
    }
} 